import torch
import torch.nn as nn
import torchvision
from torch.utils.tensorboard import writer
from vggnet import VGGNet
from load_cifar10 import train_loader, test_loader
import os
from tensorboardX import SummaryWriter


def train():
    # 是否使用GPU
    device = torch.device("cpu")
    epoch_num = 2
    lr = 0.1
    batch_size = 128  # 每次迭代训练时所使用的样本数量
    net = VGGNet().to(device)
    model_path = "./models"
    # loss
    loss_func = nn.CrossEntropyLoss()

    # optimizer
    # optimizer = torch.optim.Adam(net.parameters(), lr= lr)

    optimizer = torch.optim.SGD(net.parameters(), lr=lr,
                                momentum=0.9, weight_decay=5e-4)

    scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                step_size=10,
                                                gamma=0.9)
    writer = SummaryWriter()
    step_n = 0
    for epoch in range(epoch_num):
        print(" epoch is ", epoch)
        # 将模型开启训练模式
        net.train()  # train BN dropout

        for i, data in enumerate(train_loader):
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)

            # 向前传播： 将输入数据通过神经网络得到输出。
            outputs = net(inputs)
            # 计算损失函数值
            loss = loss_func(outputs, labels)
            optimizer.zero_grad()
            # 反向传播
            loss.backward()
            optimizer.step()

            _, pred = torch.max(outputs.data, dim=1)

            correct = pred.eq(labels.data).cpu().sum()
            print("epoch is ", epoch)
            print("train lr is ", optimizer.state_dict()["param_groups"][0]["lr"])
            print("train step", i, "loss is:", loss.item(),
                  "mini-batch correct is:", 100.0 * correct / batch_size)

            writer.add_scalar("train loss", loss.item(), global_step=0)
            writer.add_scalar("train correct",
                              100.0 * correct.item() / batch_size, global_step=0)

            im = torchvision.utils.make_grid(inputs)
            writer.add_image("train im", im, global_step=step_n)

            step_n += 1

        torch.save(net.state_dict(), "{}/{}.pth".format(model_path, epoch + 1))
        # 更新学习率
        scheduler.step()

        sum_loss = 0
        sum_correct = 0
        for i, data in enumerate(test_loader):
            # 设置为评估模式
            net.eval()
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = net(inputs)
            loss = loss_func(outputs, labels)
            _, pred = torch.max(outputs.data, dim=1)
            correct = pred.eq(labels.data).cpu().sum()

            sum_loss += loss.item()
            sum_correct += correct.item()
            im = torchvision.utils.make_grid(inputs)
            writer.add_image("test im", im, global_step=step_n)

        test_loss = sum_loss * 1.0 / len(test_loader)
        test_correct = sum_correct * 100.0 / len(test_loader) / batch_size

        writer.add_scalar("test loss", test_loss, global_step=epoch + 1)
        writer.add_scalar("test correct",
                          test_correct, global_step=epoch + 1)

        print("epoch is", epoch + 1, "loss is:", test_loss,
              "test correct is:", test_correct)

    writer.close()


if __name__ == '__main__':
    train()
